import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn import datasets

# 加载 Iris 数据集
iris = datasets.load_iris()
X = iris.data  # 特征矩阵
y = iris.target  # 目标值（类别）
feature_names = iris.feature_names  # 特征名称
target_names = iris.target_names  # 类别名称

# 创建 3D 图形
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 选择三个特征进行可视化（这里选择前三个特征）
xs = X[:, 0]  # sepal length
ys = X[:, 1]  # sepal width
zs = X[:, 2]  # petal length

# 为每个类别设置不同的颜色和标记
colors = ['r', 'g', 'b']
markers = ['o', '^', 's']

# 绘制每个类别的散点
for i, name in enumerate(target_names):
    # 筛选当前类别的数据
    mask = (y == i)
    ax.scatter(xs[mask], ys[mask], zs[mask],
               c=colors[i],
               marker=markers[i],
               label=name,
               s=50,  # 点的大小
               alpha=0.8)  # 透明度

# 添加坐标轴标签
ax.set_xlabel(feature_names[0], fontsize=12)
ax.set_ylabel(feature_names[1], fontsize=12)
ax.set_zlabel(feature_names[2], fontsize=12)

# 添加标题和图例
plt.title("Iris Dataset 3D Visualization", fontsize=14)
ax.legend(loc='best', fontsize=12)

# 调整视角（可选参数：elev=仰角, azim=方位角）
ax.view_init(elev=20, azim=-35)

plt.tight_layout()
plt.show()